package com.chatplus.application.aiprocessor.listener.impl;

import cn.bugstack.openai.executor.parameter.ChatChoice;
import cn.bugstack.openai.executor.parameter.CompletionRequest;
import cn.bugstack.openai.executor.parameter.CompletionResponse;
import cn.bugstack.openai.executor.parameter.Message;
import com.chatplus.application.aiprocessor.enumeration.MessageTypeEnum;
import com.chatplus.application.aiprocessor.listener.StreamEventSourceListener;
import com.chatplus.application.common.logging.SouthernQuietLogger;
import com.chatplus.application.common.logging.SouthernQuietLoggerFactory;
import com.chatplus.application.common.util.PlusJsonUtils;
import com.chatplus.application.enumeration.AiPlatformEnum;
import com.unfbx.chatgpt.entity.completions.Completion;
import okhttp3.sse.EventSource;
import org.apache.commons.lang3.StringUtils;
import org.jetbrains.annotations.NotNull;
import org.springframework.web.socket.WebSocketSession;

import java.util.List;

/**
 * OpenAI监听器
 */
public class CommonEventSourceListener extends StreamEventSourceListener {
    private static final SouthernQuietLogger LOGGER = SouthernQuietLoggerFactory.getLogger(CommonEventSourceListener.class);
    private String reply = "";
    // 函数参数
    private String arguments = "";
    // 函数名称
    private String name = "";

    public String getArguments() {
        return arguments;
    }

    public String getName() {
        return name;
    }
    public CommonEventSourceListener(WebSocketSession session, AiPlatformEnum platform, String prompt) {
        super(session, platform, prompt);
    }
    @Override
    public void onEvent(@NotNull EventSource eventSource, String id, String type, @NotNull String data) {
        LOGGER.message("AI返回数据").context("data", data).info();

        if ("[DONE]".equalsIgnoreCase(data)) {
            LOGGER.message("回答完成").info();
            if (StringUtils.isEmpty(name)) {
                Completion completion = Completion.builder().prompt(reply).build();
                chatAiUtilHandler.saveChatHistory(reply, completion.tokens());
                chatAiUtilHandler.sendMessage(MessageTypeEnum.WS_END.getValue());
            }
            return;
        }
        StringBuilder content = new StringBuilder();
        CompletionResponse chatCompletionResponse = PlusJsonUtils.parseObject(data, CompletionResponse.class);
        assert chatCompletionResponse != null;
        List<ChatChoice> choices = chatCompletionResponse.getChoices();
        for (ChatChoice chatChoice : choices) {
            Message delta = chatChoice.getDelta();
            if (CompletionRequest.Role.ASSISTANT.getCode().equals(delta.getRole())) continue;

            // 应答完成
            String finishReason = chatChoice.getFinishReason();
            if (StringUtils.isNoneBlank(finishReason) && "stop".equalsIgnoreCase(finishReason)) {
                return;
            }
            content.append(delta.getContent());
        }
        if (StringUtils.isNotEmpty(content.toString())) {
            chatAiUtilHandler.sendMessage(content.toString());
            reply += content.toString();
        }
    }
}
